/**
 * 
 */
package edu.umd.clip.lm.storage;

import java.io.*;

import edu.umd.clip.lm.factors.FactorTupleDescription;
import edu.umd.clip.lm.model.Experiment;
import edu.umd.clip.lm.model.OnDiskCompactProbTree;
import com.sleepycat.je.*;
import com.sleepycat.bind.tuple.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class BDBProbTreeStorage extends AbstractProbTreeStorage {
	private final Environment env;
	private Database dbs[];
	private final KeyBinding keyBindings[];
	private final static ProbTreeBinding probTreeBinding = new ProbTreeBinding();
	private static final int MAX_LOG_FILE_SIZE = 500 * 1024 * 1024;
	private final static int MAX_MODELS = 5;
	public BDBProbTreeStorage(Environment env) throws DatabaseException {
		this.env = env;
		this.keyBindings = new KeyBinding[MAX_MODELS];
		this.dbs = new Database[MAX_MODELS];
	}
	
	public static Environment createEnvironment(String path, boolean forWriting) throws DatabaseException {
		EnvironmentConfig envConf = new EnvironmentConfig();
		envConf.setAllowCreate(forWriting);
		//envConf.setSharedCache(true);
		envConf.setTransactional(false);
		envConf.setReadOnly(!forWriting);
		envConf.setLocking(forWriting);
		envConf.setCachePercent(20);
		envConf.setConfigParam("je.log.fileMax", Integer.toString(MAX_LOG_FILE_SIZE));
		
		File dir = new File(path);
		if (!dir.isDirectory()) {
			dir.mkdirs();
		}
		Environment env = new Environment(dir, envConf);
		return env;
	}
	
	public void open(String dbname, int modelId, boolean forWriting) throws DatabaseException {
		DatabaseConfig dbConf = new DatabaseConfig();
		dbConf.setReadOnly(!forWriting);
		dbConf.setAllowCreate(forWriting);
		dbConf.setTransactional(false);
		dbConf.setDeferredWrite(true);
		
		if (forWriting) {
			try {
				env.truncateDatabase(null, dbname, false);
			} catch(DatabaseNotFoundException e) {
			}
		}
		
		dbs[modelId] = env.openDatabase(null, dbname, dbConf);
		keyBindings[modelId] = new KeyBinding(modelId);
	}


	public void sync(int modelId) throws DatabaseException {
		dbs[modelId].sync();
	}
	
	public void close(int modelId) throws DatabaseException {
		dbs[modelId].close();
	}
	
	public static class KeyBinding extends TupleBinding<Key> {
		private final int modelid;
		
		public KeyBinding(int modelid) {
			this.modelid = modelid;
		}

		/* (non-Javadoc)
		 * @see com.sleepycat.bind.tuple.TupleBinding#entryToObject(com.sleepycat.bind.tuple.TupleInput)
		 */
		@Override
		public Key entryToObject(TupleInput ti) {
			FactorTupleDescription desc = Experiment.getInstance().getTupleDescription();
			byte bytes[] = new byte[desc.getCompactOvertFactorsBytes()];
			ti.readFast(bytes);
			long factorBits = desc.unpackOvertFactors(bytes);
			int clusterid = ti.readInt();
			return new Key(modelid, clusterid, factorBits);
		}

		/* (non-Javadoc)
		 * @see com.sleepycat.bind.tuple.TupleBinding#objectToEntry(java.lang.Object, com.sleepycat.bind.tuple.TupleOutput)
		 */
		@Override
		public void objectToEntry(Key key, TupleOutput to) {
			byte bytes[] = Experiment.getInstance().getTupleDescription().packOvertFactors(key.factorBits);
			to.writeFast(bytes);
			to.writeInt(key.clusterid);
		}
	}
	
	public static class ProbTreeBinding extends TupleBinding<OnDiskCompactProbTree> {
		/* (non-Javadoc)
		 * @see com.sleepycat.bind.tuple.TupleBinding#entryToObject(com.sleepycat.bind.tuple.TupleInput)
		 */
		@Override
		public OnDiskCompactProbTree entryToObject(TupleInput ti) {
			OnDiskCompactProbTree tree = new OnDiskCompactProbTree(ti);
			return tree;
		}

		/* (non-Javadoc)
		 * @see com.sleepycat.bind.tuple.TupleBinding#objectToEntry(java.lang.Object, com.sleepycat.bind.tuple.TupleOutput)
		 */
		@Override
		public void objectToEntry(OnDiskCompactProbTree tree, TupleOutput to) {
			tree.writeTuple(to);
		}
	}
	
	@Override
	public OnDiskCompactProbTree getProbTree(Key key) {
		DatabaseEntry keyData = new DatabaseEntry();
		DatabaseEntry probTreeData = new DatabaseEntry();
		keyBindings[key.modelid].objectToEntry(key, keyData);
		
		try {
			OperationStatus status = dbs[key.modelid].get(null, keyData, probTreeData, null);
			if (status == OperationStatus.SUCCESS) {
				OnDiskCompactProbTree tree = probTreeBinding.entryToObject(probTreeData);
				return tree;
			} else if (status != OperationStatus.NOTFOUND) {
				System.err.printf("DB reading failure: %s\n" + status.toString());
			}
		} catch (DatabaseException e) {
			e.printStackTrace();
		}
		return null;
	}


	@Override
	public void setProbTree(Key key, OnDiskCompactProbTree probTree) {
		DatabaseEntry keyData = new DatabaseEntry();
		DatabaseEntry probTreeData = new DatabaseEntry();
		keyBindings[key.modelid].objectToEntry(key, keyData);
		probTreeBinding.objectToEntry(probTree, probTreeData);
		
		try {
			OperationStatus status = dbs[key.modelid].put(null, keyData, probTreeData);
			if (status != OperationStatus.SUCCESS) {
				System.err.printf("DB writing failure: %s\n" + status.toString());				
			}
		} catch (DatabaseException e) {
			e.printStackTrace();
		}
	}
	
	public void removeProbTree(int modelid, int clusterid, long factorBits) {
		removeProbTree(new Key(modelid, clusterid, factorBits));
	}
	
	public void removeProbTree(Key key) {
		DatabaseEntry keyData = new DatabaseEntry();
		keyBindings[key.modelid].objectToEntry(key, keyData);
		try {
			OperationStatus status = dbs[key.modelid].delete(null, keyData);
			if (status != OperationStatus.SUCCESS) {
				System.err.printf("DB record deletion failure: %s\n", status.toString());
			}
		} catch (DatabaseException e) {
			e.printStackTrace();
		}
	}
	
	public EnvironmentStats getEnvStats() throws DatabaseException {
		return env.getStats(null);
	}
	
	public DatabaseStats getDbStats(int modelId) throws DatabaseException {
		return dbs[modelId].getStats(null);
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AbstractProbTreeStorage#getStats()
	 */
	@Override
	public String getStats() {
		try {
			StringBuilder sb = new StringBuilder(getEnvStats().toString());
			for(Database db : dbs) {
				if (db != null) {
					sb.append('\n');
					sb.append(db.getStats(null));
				}
			}
			return  sb.toString();
		} catch (DatabaseException e) {
			e.printStackTrace();
		}
		return "";
	}

	@Override
	protected void requestProbTree(Key key) {
		// noop
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.storage.AsyncProducer#request(java.lang.Object, edu.umd.clip.lm.storage.AsyncConsumer)
	 */
	@Override
	public void request(Key key, AsyncConsumer<Key, OnDiskCompactProbTree> callback) {
		OnDiskCompactProbTree probTree = getProbTree(key);
		callback.receivedData(key, probTree);
	}

	@Override
	public void closeAll() {
		for(int i=0; i<dbs.length; ++i) {
			if (dbs[i] != null) {
				close(i);
			}
		}
	}
	
}
